import type { GOOGLE_GENAI_INSTRUMENTED_METHODS } from './constants';

export interface GoogleGenAIOptions {
  /**
   * Enable or disable input recording.
   */
  recordInputs?: boolean;
  /**
   * Enable or disable output recording.
   */
  recordOutputs?: boolean;
}

/**
 * Google GenAI Content Part
 * @see https://ai.google.dev/api/rest/v1/Content#Part
 * @see https://github.com/googleapis/js-genai/blob/v1.19.0/src/types.ts#L1061
 *
 */
export type ContentPart = {
  /** Metadata for a given video. */
  videoMetadata?: unknown;
  /** Indicates if the part is thought from the model. */
  thought?: boolean;
  /** Optional. Inlined bytes data. */
  inlineData?: Blob;
  /** Optional. URI based data. */
  fileData?: unknown;
  /** An opaque signature for the thought so it can be reused in subsequent requests.
   * @remarks Encoded as base64 string. */
  thoughtSignature?: string;
  /** A predicted [FunctionCall] returned from the model that contains a string
       representing the [FunctionDeclaration.name] and a structured JSON object
       containing the parameters and their values. */
  functionCall?: {
    /** The unique id of the function call. If populated, the client to execute the
   `function_call` and return the response with the matching `id`. */
    id?: string;
    /** Optional. The function parameters and values in JSON object format. See [FunctionDeclaration.parameters] for parameter details. */
    args?: Record<string, unknown>;
    /** Required. The name of the function to call. Matches [FunctionDeclaration.name]. */
    name?: string;
  };
  /** Optional. Result of executing the [ExecutableCode]. */
  codeExecutionResult?: unknown;
  /** Optional. Code generated by the model that is meant to be executed. */
  executableCode?: unknown;
  /** Optional. The result output of a [FunctionCall] that contains a string representing the [FunctionDeclaration.name] and a structured JSON object containing any output from the function call. It is used as context to the model. */
  functionResponse?: unknown;
  /** Optional. Text part (can be code). */
  text?: string;
};

/**
 * Google GenAI Content
 * @see https://ai.google.dev/api/rest/v1/Content
 */
type Content = {
  /** List of parts that constitute a single message.
   * Each part may have a different IANA MIME type. */
  parts?: ContentPart[];
  /** Optional. The producer of the content. Must be either 'user' or
   * 'model'. Useful to set for multi-turn conversations, otherwise can be
   * empty. If role is not specified, SDK will determine the role.
   */
  role?: string;
};

type MediaModality = 'MODALITY_UNSPECIFIED' | 'TEXT' | 'IMAGE' | 'VIDEO' | 'AUDIO' | 'DOCUMENT';

/**
 * Google GenAI Modality Token Count
 * @see https://ai.google.dev/api/rest/v1/ModalityTokenCount
 */
type ModalityTokenCount = {
  /** The modality associated with this token count. */
  modality?: MediaModality;
  /** Number of tokens. */
  tokenCount?: number;
};

/**
 * Google GenAI Usage Metadata
 * @see https://ai.google.dev/api/rest/v1/GenerateContentResponse#UsageMetadata
 */
type GenerateContentResponseUsageMetadata = {
  [key: string]: unknown;
  /** Output only. List of modalities of the cached content in the request input. */
  cacheTokensDetails?: ModalityTokenCount[];
  /** Output only. Number of tokens in the cached part in the input (the cached content). */
  cachedContentTokenCount?: number;
  /** Number of tokens in the response(s). */
  candidatesTokenCount?: number;
  /** Output only. List of modalities that were returned in the response. */
  candidatesTokensDetails?: ModalityTokenCount[];
  /** Number of tokens in the request. When `cached_content` is set, this is still the total effective prompt size meaning this includes the number of tokens in the cached content. */
  promptTokenCount?: number;
  /** Output only. List of modalities that were processed in the request input. */
  promptTokensDetails?: ModalityTokenCount[];
  /** Output only. Number of tokens present in thoughts output. */
  thoughtsTokenCount?: number;
  /** Output only. Number of tokens present in tool-use prompt(s). */
  toolUsePromptTokenCount?: number;
  /** Output only. List of modalities that were processed for tool-use request inputs. */
  toolUsePromptTokensDetails?: ModalityTokenCount[];
  /** Total token count for prompt, response candidates, and tool-use prompts (if present). */
  totalTokenCount?: number;
};

/**
 * Google GenAI Candidate
 * @see https://ai.google.dev/api/rest/v1/Candidate
 * https://github.com/googleapis/js-genai/blob/v1.19.0/src/types.ts#L2237
 */
export type Candidate = {
  [key: string]: unknown;
  /**
   * Contains the multi-part content of the response.
   */
  content?: Content;
  /**
   * The reason why the model stopped generating tokens.
   * If empty, the model has not stopped generating the tokens.
   */
  finishReason?: string;
  /**
   * Number of tokens for this candidate.
   */
  tokenCount?: number;
  /**
   * The index of the candidate.
   */
  index?: number;
};

/**
 * Google GenAI Generate Content Response
 * @see https://ai.google.dev/api/rest/v1/GenerateContentResponse
 */
type GenerateContentResponse = {
  [key: string]: unknown;
  /** Response variations returned by the model. */
  candidates?: Candidate[];
  /** Timestamp when the request is made to the server. */
  automaticFunctionCallingHistory?: Content[];
  /** Output only. The model version used to generate the response. */
  modelVersion?: string;
  /** Output only. Content filter results for a prompt sent in the request. Note: Sent only in the first stream chunk. Only happens when no candidates were generated due to content violations. */
  promptFeedback?: Record<string, unknown>;
  /** Output only. response_id is used to identify each response. It is the encoding of the event_id. */
  responseId?: string;
  /** Usage metadata about the response(s). */
  usageMetadata?: GenerateContentResponseUsageMetadata;
};

/**
 * Basic interface for Google GenAI client with only the instrumented methods
 * This provides type safety while being generic enough to work with different client implementations
 */
export interface GoogleGenAIClient {
  models: {
    generateContent: (...args: unknown[]) => Promise<GenerateContentResponse>;
    // https://googleapis.github.io/js-genai/release_docs/classes/models.Models.html#generatecontentstream
    // eslint-disable-next-line @typescript-eslint/no-explicit-any
    generateContentStream: (...args: unknown[]) => Promise<AsyncGenerator<GenerateContentResponse, any, unknown>>;
  };
  chats: {
    create: (...args: unknown[]) => GoogleGenAIChat;
  };
}

/**
 * Google GenAI Chat interface for chat instances created via chats.create()
 */
export interface GoogleGenAIChat {
  sendMessage: (...args: unknown[]) => Promise<GenerateContentResponse>;
  // https://googleapis.github.io/js-genai/release_docs/classes/chats.Chat.html#sendmessagestream
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  sendMessageStream: (...args: unknown[]) => Promise<AsyncGenerator<GenerateContentResponse, any, unknown>>;
}

export type GoogleGenAIIstrumentedMethod = (typeof GOOGLE_GENAI_INSTRUMENTED_METHODS)[number];

// Export the response type for use in instrumentation
export type GoogleGenAIResponse = GenerateContentResponse;
